import yt_dlp
from typing import Dict, List, Optional, Tuple, Callable
from .base import BaseDownloader
import os

class YouTubeDownloader(BaseDownloader):
    def __init__(self, config):
        super().__init__(config)
        self.ydl_opts = {
            'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best',
            'outtmpl': '%(title)s.%(ext)s',
            'quiet': True,
            'no_warnings': True,
            'extract_flat': False,
            'retries': 3,
            'fragment_retries': 3,
            'retry_sleep': 5,
            'fragment_retry_sleep': lambda n: 5 * (n + 1),
            'socket_timeout': 30,
            'headers': self.headers
        }
    
    async def extract_video_info(self, url: str) -> Dict:
        """提取视频信息"""
        try:
            with yt_dlp.YoutubeDL(self.ydl_opts) as ydl:
                info = ydl.extract_info(url, download=False)
                
                if not info:
                    raise Exception("无法获取视频信息")
                
                # 获取可用的格式
                formats = []
                for f in info.get('formats', []):
                    if f.get('ext') == 'mp4':
                        formats.append({
                            'format_id': f.get('format_id', ''),
                            'ext': f.get('ext', ''),
                            'format_note': f.get('format_note', ''),
                            'resolution': f.get('resolution', ''),
                            'filesize': f.get('filesize', 0),
                            'tbr': f.get('tbr', 0)
                        })
                
                # 按分辨率排序
                formats.sort(key=lambda x: (
                    int(x['resolution'].split('x')[0]) if x['resolution'] else 0,
                    x['tbr']
                ), reverse=True)
                
                return self._format_video_info(
                    title=info.get('title', ''),
                    formats=formats,
                    cover_url=info.get('thumbnail', '')
                )
                
        except yt_dlp.utils.DownloadError as e:
            if "Video unavailable" in str(e):
                raise Exception("视频不可用，可能已被删除")
            elif "Sign in to confirm your age" in str(e):
                raise Exception("需要登录才能观看此视频")
            else:
                raise Exception(f"获取视频信息失败: {str(e)}")
        except Exception as e:
            raise Exception(f"解析视频失败: {str(e)}")
    
    async def download_video(self, url: str, format_id: str, 
                      progress_callback: Optional[Callable] = None) -> Tuple[bool, str]:
        """下载视频"""
        try:
            # 更新下载选项
            self.ydl_opts.update({
                'format': format_id,
                'outtmpl': os.path.join(self.config.download_dir, '%(title)s.%(ext)s'),
                'progress_hooks': [lambda d: self._handle_progress(d, progress_callback)]
            })
            
            # 开始下载
            with yt_dlp.YoutubeDL(self.ydl_opts) as ydl:
                info = ydl.extract_info(url, download=True)
                if not info:
                    return False, "下载失败：无法获取视频信息"
                
                # 获取下载文件路径
                save_path = os.path.join(
                    self.config.download_dir,
                    self._clean_filename(info['title']) + '.' + info['ext']
                )
                
                if os.path.exists(save_path) and os.path.getsize(save_path) > 0:
                    return True, save_path
                else:
                    return False, "下载失败：文件未生成"
                    
        except Exception as e:
            return False, f"下载失败：{str(e)}"
    
    def _handle_progress(self, d: Dict, callback: Optional[Callable] = None):
        """处理下载进度"""
        if callback and d['status'] == 'downloading':
            try:
                total = d.get('total_bytes', 0)
                downloaded = d.get('downloaded_bytes', 0)
                if total > 0:
                    progress = (downloaded / total) * 100
                    callback(progress)
            except Exception:
                pass 